from enum import Enum
from typing import Dict
from torch.utils.data import DataLoader


class DatasetType(Enum):
    CITYSCAPES = 'cityscapes'
    ANIMALS = 'animals'
    Radar = 'radar'

    @classmethod
    def from_str(cls, label: str) -> "DatasetType":
        if label in cls.__members__:
            return cls[label]
        
        for member in cls:
            if member.value.lower() == label.lower():
                return member
        raise ValueError(f"Unknown DatasetType: {label!r}. "
                         f"Valid names: {list(cls.__members__.keys())}, "
                         f"values: {[m.value for m in cls]}")


def get_dataset(cfg: Dict[str, dict]):
    dataset_type = cfg.get('dataset')
    if isinstance(dataset_type, str):
        dataset_type = DatasetType.from_str(dataset_type)
    batch_size = cfg.get('batch_size')

    if dataset_type == DatasetType.CITYSCAPES:
        from .cityscape.generator import CityScapes
        train_dataset = CityScapes('train', cfg)
        val_dataset = CityScapes('val', cfg)
    elif dataset_type == DatasetType.ANIMALS:
        from .animals.generator import Animals
        train_dataset = Animals('train', cfg)
        val_dataset = Animals('val', cfg)
    elif dataset_type == DatasetType.Radar:
        from .radar.generator import Radar
        train_dataset = Radar('train', cfg)
        val_dataset = Radar('val', cfg)
    
    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=4, drop_last=True)
    val_loader = DataLoader(val_dataset, batch_size=1, shuffle=False, num_workers=4)
    return train_loader, val_loader
